import logging
from termcolor import colored

# 自定义日志处理器
class ColoredFormatter(logging.Formatter):
    def format(self, record):
        # 指定不同级别的颜色
        if record.levelno == logging.INFO:
            color = 'cyan'
        elif record.levelno == logging.WARNING:
            color = 'yellow'
        elif record.levelno == logging.ERROR:
            color = 'red'
        elif record.levelno == logging.DEBUG:
            color = 'blue'
        else:
            color = 'white'
        
        # 格式化消息并应用颜色
        message = super().format(record)
        return colored(message, color)

# 配置日志记录器
logger = logging.getLogger(__name__)
handler = logging.StreamHandler()

# 设置处理器使用自定义格式化
formatter = ColoredFormatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
handler.setFormatter(formatter)

# 添加处理器到记录器
logger.addHandler(handler)
logger.setLevel(logging.DEBUG)

def _create_realtime_callback(task_name: str, host_name: str):
    """
    为流式命令创建实时输出回调函数
    
    Args:
        task_name: 任务名称
        host_name: 主机名称
        
    Returns:
        function: 回调函数
    """
    async def callback(line: str, result):
        """实时输出回调"""
        # 过滤掉过长的行，避免显示混乱
        display_line = line[:100] + "..." if len(line) > 100 else line
        
        # 颜色映射
        colors = ['red', 'green', 'yellow', 'blue', 'magenta', 'cyan']
        
        # 根据主机名称选择颜色
        color_index = hash(host_name) % len(colors)
        color = colors[color_index]

        # 实时显示，使用 termcolor 输出颜色
        colored_host_name = colored(host_name, color)
        logger.info(f"{colored_host_name} | {display_line}")

    return callback

# 测试示例
if __name__ == "__main__":
    callback_function = _create_realtime_callback(None, "personal-server-1")
    
    # 模拟实时输出
    example_line = "This is a sample log output that might be long and needs to be truncated."
    
    import asyncio
    asyncio.run(callback_function(example_line, None))